package trie;

import java.util.TreeMap;

/**
 * 
 * @Title: CandyTrie.java 
 * @Package trie 
 * @Description: 自定义字典树
 * @author CandyWall   
 * @date 2021年1月10日 下午5:09:23 
 * @version V1.0
 */
public class CandyTrie {
    private class Node {
        public boolean isWord;
        public TreeMap<Character, Node> next;
        public Node(boolean isWord) {
            this.isWord = isWord;
            next = new TreeMap<>();
        }
        
        public Node() {
            this(false);
        }
    }
    
    private Node root;
    private int size;
    
    public CandyTrie() {
        this.root = new Node();
        this.size = 0;
    }
    
    public int getSize() {
        return size;
    }
    
    // 向Trie中添加一个单词Word（非递归写法）
    public void add(String word) {
        Node cur = root;
        for(int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if(!cur.next.containsKey(c)) {
                cur.next.put(c, new Node());
            }
            cur = cur.next.get(c);
        }
        // 确保第一次添加这个单词，数量才加1
        if(!cur.isWord) {
            cur.isWord = true;
            size++;
        }
    }
    
    // 向Trie中添加一个单词word（递归写法）
    /*public void add(String word) {
        for(int i = 0; i < word.length(); i++) {
            
        }
    }
    private Node add(Node node, Character c) {
        if(node.next == null) {
           return node
        }
    }*/
    
    // 查询单词Word是否在Trie中
    public boolean contains(String word) {
        Node cur = root;
        for(int i = 0; i < word.length(); i++) {
            // 如果当前节点的下一个节点中不包含当前字母，就表明单词word不在Trie中
            char c = word.charAt(i);
            if(!cur.next.containsKey(c)) {
                return false;
            }
            cur = cur.next.get(c);
        }
        // 不能直接返回return true; 因为返回的单词可能只是其他更长的单词的前缀，并没有存进来
        return cur.isWord;       
    }
    
    // 查询是否在Trie中有单词以prefix为前缀
    public boolean isPrefix(String prefix) {
        Node cur = root;
        for(int i = 0; i < prefix.length(); i++) {
            // 如果当前节点的下一个节点中不包含当前字母，就表明单词word不在Trie中
            char c = prefix.charAt(i);
            if(!cur.next.containsKey(c)) {
                return false;
            }
            cur = cur.next.get(c);
        }
        return true;
    }
    
    public boolean isEmpty() {
        return size == 0;
    }
}
